/*
 * Copyright (C) 2013 Aurora Cain
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package tavernaPBS;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.CommandLineParser;
import org.apache.commons.cli.GnuParser;
import org.apache.commons.cli.HelpFormatter;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.OptionBuilder;
import org.apache.commons.cli.ParseException;
import org.apache.commons.cli.UnrecognizedOptionException;
import tavernaPBS.Job;
import tavernaPBS.QSubException;
import tavernaPBS.TorquePBS;

/**
 *
 * @author aacain
 */
public class BlastPBSRun {

    private static String endcoding;
    private String database;
    private List<String> modules;
    private String outDirectory;
    private String blastPath;
    private HashMap<String, String> jobToFileMap = new HashMap<String, String>();

    /**
     *
     * @param database - full path and name of database (required)
     * @param outDirectory - the directory in which files will be output for the
     * blast run (required)
     * @param modules - list of environment modules to load (optional)
     * @param blastPath - the directory of the blast program (optional)
     */
    public BlastPBSRun(String database, String outDirectory, List<String> modules, String blastPath) {
        this.database = database;
        this.modules = modules;
        if (outDirectory != null && !outDirectory.isEmpty()) {
            if (outDirectory.charAt(outDirectory.length() - 1) != '/') {
                outDirectory = outDirectory + '/';
            }
        }
        this.outDirectory = outDirectory;
        this.blastPath = blastPath;
    }

    /**
     *
     * @param inPBS -
     * @param program
     * @param queryFileNames
     * @param additionalParams
     * @param runAfterIds
     * @return
     * @throws IOException
     * @throws Exception
     */
    public List<String> runBlast(TorquePBS inPBS, String program, List<String> queryFileNames, List<String> additionalParams, List<String> runAfterIds) throws IOException, Exception {
        List<String> currentJobIds = new LinkedList<String>();
        inPBS.closeRemoteConnection();
        for (String queryFile : queryFileNames) {
            System.out.println("Submitting job for " + queryFile);
            currentJobIds.add(runBlast(inPBS, program, queryFile, additionalParams, runAfterIds));
        }

        return currentJobIds;
    }

    public String runBlast(TorquePBS inPBS, String program, String queryFileName, List<String> additionalParams, List<String> runAfterIds) throws IOException, Exception {
        ArrayList<String> addCommands = new ArrayList<String>();
        StringBuilder path = new StringBuilder();
        if (outDirectory != null && !outDirectory.isEmpty()) {
            path.append(outDirectory);
            addCommands.add("DIRECTORY=" + path.toString());
            addCommands.add("if [ ! -d \"$DIRECTORY\" ]; then");
            addCommands.add("mkdir -p \"$DIRECTORY\"");
            addCommands.add("fi");
        }
        if (modules != null) {
            for (String module : modules) {
                addCommands.add("module load " + module);
            }
        }

        StringBuilder outFileName = new StringBuilder(path);
        int extIndex = queryFileName.lastIndexOf('.');
        int pathIndex = queryFileName.lastIndexOf('/');
        outFileName.append(queryFileName.subSequence(pathIndex >= 0 ? pathIndex + 1 : 0,
                extIndex >= 0 ? extIndex : queryFileName.length()));
        outFileName.append(".out");

        StringBuilder blast = new StringBuilder();
        if (blastPath != null && !blastPath.isEmpty()) {
            blast.append(blastPath);
            if (!blastPath.endsWith("/")) {
                blast.append('/');
            }
        }
        blast.append(program).append(" -query ")
                .append(queryFileName)
                .append(" -db ").append(database)
                .append(" -out ").append(outFileName);
        if (additionalParams != null) {
            for (String param : additionalParams) {
                blast.append(' ').append(param);
            }
        }
        addCommands.add(blast.toString());

        //create job and run
        Job job = null;
        String jobId = null;
        for (String command : addCommands) {
            if (job == null) {
                job = new Job(command);
            } else {
                job.addCommand(command);
            }
        }
        if (runAfterIds == null) {
            jobId = inPBS.doJob(job);
        } else {
            jobId = inPBS.doJob(job, runAfterIds);
        }
        this.jobToFileMap.put(jobId, outFileName.toString());

        inPBS.closeRemoteConnection();
        return jobId;
    }

    public String createDatabase(TorquePBS inPBS, List<String> fileNames, String databaseType, List<String> additionalParams, List<String> runAfterIds) throws IOException, QSubException {
        ArrayList<String> addCommands = new ArrayList<String>();
        int lastIndexOf = database.lastIndexOf('/');
        if (lastIndexOf >= 0) {
            String dbDirectory = database.substring(0, lastIndexOf);
            addCommands.add("DIRECTORY=" + dbDirectory);
            addCommands.add("if [ ! -d \"$DIRECTORY\" ]; then");
            addCommands.add("mkdir -p \"$DIRECTORY\"");
            addCommands.add("fi");
        }

        if (modules != null) {
            for (String module : modules) {
                addCommands.add("module load " + module);
            }
        }
        StringBuilder makeBlastDb = new StringBuilder();
        if (blastPath != null && !blastPath.isEmpty()) {
            makeBlastDb.append(blastPath);
            if (!blastPath.endsWith("/")) {
                makeBlastDb.append('/');
            }
        }
        makeBlastDb.append("makeblastdb -in \"");
        for (String fileName : fileNames) {
            makeBlastDb.append(fileName).append(' ');
        }
        makeBlastDb.replace(makeBlastDb.length() - 1, makeBlastDb.length(), "\" ");
        makeBlastDb.append(" -dbtype ").append(databaseType).append(" -out ").append(database);
        if (additionalParams != null) {
            for (String param : additionalParams) {
                makeBlastDb.append(' ').append(param);
            }
        }
        addCommands.add(makeBlastDb.toString());

        //create the job and run
        Job job = null;
        String jobId = null;
        for (String command : addCommands) {
            if (job == null) {
                job = new Job(command);
            } else {
                job.addCommand(command);
            }
        }
        if (runAfterIds == null) {
            jobId = inPBS.doJob(job);
        } else {
            jobId = inPBS.doJob(job, runAfterIds);
        }

        inPBS.closeRemoteConnection();

        return jobId;
    }

    public static List<String> splitFile(String localFileName, String outDirectory, int numFiles) throws FileNotFoundException, IOException {
        List<String> fileNames = new ArrayList<String>(numFiles);
        StringBuilder outFilePrefix = new StringBuilder();
        int extIndex = localFileName.lastIndexOf('.');
        int pathIndex = localFileName.lastIndexOf('/');
        outFilePrefix.append(localFileName.subSequence(pathIndex >= 0 ? pathIndex + 1 : 0,
                extIndex >= 0 ? extIndex : localFileName.length()));
        outFilePrefix.append('_');
        CharSequence ext = localFileName.subSequence(extIndex >= 0 ? extIndex + 1 : 0, localFileName.length());

        File file = new File(localFileName);
        long length = file.length() / numFiles;
        InputStreamReader r = new InputStreamReader(new FileInputStream(file));
        String encoding = r.getEncoding();
        BufferedReader br = new BufferedReader(r);
        String line = null;
        BufferedWriter bw = null;
        try {
            Integer fileNumber = 0;
            long writtenBytes = 0;
            while ((line = br.readLine()) != null) {
                int bytes = line.getBytes(endcoding).length;
                if (line.startsWith(">")) {
                    //check to see if we should use new file
                    if (writtenBytes > length) {
                        bw.close();
                        StringBuilder fileName = new StringBuilder(outFilePrefix);
                        fileName.append(fileNumber.toString()).append('.').append(ext);
                        File currentFile = new File(fileName.toString());
                        currentFile.createNewFile();
                        bw = new BufferedWriter(new OutputStreamWriter(new FileOutputStream(currentFile), encoding));
                        fileNames.add(fileName.toString());
                    }
                }
                bw.write(line);
                bw.newLine();
                writtenBytes = writtenBytes + bytes;
            }
        } finally {
            try {
                br.close();
            } catch (IOException ex) {
                Logger.getLogger(BlastPBSRun.class.getName()).log(Level.SEVERE, null, ex);
            }
            if (bw != null) {
                try {
                    bw.close();
                } catch (IOException ex) {
                    Logger.getLogger(BlastPBSRun.class.getName()).log(Level.SEVERE, null, ex);
                }
            }
        }
        return fileNames;
    }

    public static void testJobSubmission(TorquePBS inPBS, String echo) throws IOException, Exception {
        System.out.println("Submitting job " + echo);
        Job job = new Job("echo " + echo);
        inPBS.doJob(job);
        inPBS.closeRemoteConnection();
    }

    public static void main(String[] args) throws IOException, Exception {
        System.err.println("Running BlastPBSRun");
        Option help = new Option("h", "help", false, "print this message");
        Option directory = OptionBuilder.withArgName("directory")
                .hasArg()
                .withDescription("directory of fasta files (required)")
                .withLongOpt("in")
                .isRequired()
                .create("i");
        Option database = OptionBuilder.withArgName("database")
                .hasArg()
                .withDescription("full path of the database (required)")
                .isRequired()
                .withLongOpt("databaseName")
                .create("d");
        Option outDir = OptionBuilder.withArgName("out")
                .hasArg()
                .withDescription("result output directory (required)")
                .isRequired()
                .withLongOpt("outDirectory")
                .create("o");
        Option makeDb = new Option("c", "createDb", false, "create the database");
        Option mod = OptionBuilder.withArgName("modules")
                .hasArgs()
                .withValueSeparator(' ')
                .withDescription("PBS modules to load - seperated by space (optional)")
                .withLongOpt("modules")
                .create("m");
        Option type = OptionBuilder.withArgName("program")
                .hasArg()
                .isRequired()
                .withDescription("blast program to run (blastp, blastn, ...) (required)")
                .withLongOpt("program")
                .create("p");
        Option blastParams = OptionBuilder.withArgName("blastParams")
                .hasArgs()
                .withDescription("blast program parameters (optional)")
                .withLongOpt("additionalParams")
                .create("a");
        Option bin = OptionBuilder.withArgName("bin")
                .hasArg()
                .withDescription("blast bin (optional)")
                .withLongOpt("bin")
                .create("b");
        Option extension = OptionBuilder.withArgName("ext")
                .hasArg()
                .withDescription("fasta file extension - default = fasta (optional)")
                .withLongOpt("ext")
                .create("e");

        Options options1 = new Options();
        options1.addOption(help);

        Options options2 = new Options();
        options2.addOption(directory);
        options2.addOption(database);
        options2.addOption(outDir);
        options2.addOption(makeDb);
        options2.addOption(mod);
        options2.addOption(type);
        options2.addOption(blastParams);
        options2.addOption(bin);
        options2.addOption(extension);


        CommandLineParser parser = new GnuParser();
        try {
            CommandLine c1 = parser.parse(options1, args);
            if (c1.getOptions() != null && c1.getOptions().length > 0) {
                if (c1.hasOption("help")) {
                    HelpFormatter formatter = new HelpFormatter();
                    formatter.printHelp("blastPBS", "", options2, "Use '-h' or '--help' to print this message", true);
                }
            }else{
                run(args, options2);
            }
        } catch (ParseException ex1) {
            run(args, options2);
        }
    }

    private static void run(String[] args, Options options2) throws IOException, QSubException {
        try {
            CommandLineParser parser = new GnuParser();
            CommandLine c1 = parser.parse(options2, args);
            List<String> mods = null;
            if (c1.hasOption("m")) {
                mods = new ArrayList(Arrays.asList(c1.getOptionValues("modules")));
            }
            BlastPBSRun run = new BlastPBSRun(c1.getOptionValue("d"), c1.getOptionValue("o"), mods, c1.getOptionValue("b"));
            String dir = c1.getOptionValue("in");
            String prog = c1.getOptionValue("program");
            String ext = c1.getOptionValue("ext");
            if (ext == null) {
                ext = "fasta";
            }
            //get the files
            TorquePBS pbs = new TorquePBS();
            pbs.startLog();
            List<String> remoteFiles = new ArrayList(Arrays.asList(pbs.getRemoteFiles(dir, ext)));
            if (c1.hasOption("c")) {
                run.createDatabase(pbs, remoteFiles, ext, null, null);
            }
        } catch (UnrecognizedOptionException ex1) {
            System.err.println("Unexpected exception:" + ex1.getMessage());
            HelpFormatter formatter = new HelpFormatter();
            formatter.printHelp("blastPBS", "", options2, "Use '-h' or '--help' to print this message", true);
        } catch (ParseException ex1) {
            System.err.println("Unexpected exception:" + ex1.getMessage());
            HelpFormatter formatter = new HelpFormatter();
            formatter.printHelp("blastPBS", "", options2, "Use '-h' or '--help' to print this message", true);
        }
    }
}
